#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
联邦学习配置管理

功能：
1. 参数配置管理
2. 客户端注册
3. 训练调度
4. 日志记录
"""

import yaml
from typing import Dict, List, Any
import os
import logging
from datetime import datetime

class FederatedConfig:
    """联邦学习配置管理器"""
    
    def __init__(self, config_path: str = None):
        """
        初始化配置管理器
        
        参数:
            config_path: 配置文件路径
        """
        self.default_config = {
            'federated': {
                'rounds': 100,
                'clients_per_round': 5,
                'aggregation': 'weighted_average',
                'model': {
                    'type': 'IndustrialFederatedModel',
                    'input_size': 10,
                    'output_size': 5
                },
                'optimizer': {
                    'type': 'SGD',
                    'lr': 0.01,
                    'momentum': 0.9
                },
                'privacy': {
                    'dp_enabled': True,
                    'dp_sigma': 0.1,
                    'clip_norm': 1.0
                },
                'security': {
                    'byzantine_threshold': 0.3,
                    'encryption': 'rsa'
                }
            }
        }
        
        self.config = self.default_config
        if config_path and os.path.exists(config_path):
            self.load_config(config_path)
            
        self._setup_logging()
        self.client_registry = {}
        
    def load_config(self, path: str):
        """从YAML文件加载配置"""
        with open(path, 'r') as f:
            self.config = yaml.safe_load(f) or self.default_config
            
    def save_config(self, path: str):
        """保存配置到YAML文件"""
        with open(path, 'w') as f:
            yaml.dump(self.config, f)
            
    def register_client(self, client_id: str, metadata: Dict):
        """注册客户端"""
        self.client_registry[client_id] = {
            'metadata': metadata,
            'last_active': datetime.now(),
            'performance': []
        }
        
    def update_client_performance(self, client_id: str, metrics: Dict):
        """更新客户端性能指标"""
        if client_id in self.client_registry:
            self.client_registry[client_id]['performance'].append(metrics)
            self.client_registry[client_id]['last_active'] = datetime.now()
            
    def get_client_selection(self, round_num: int) -> List[str]:
        """获取本轮训练的客户端选择"""
        # 简单实现：随机选择
        all_clients = list(self.client_registry.keys())
        np.random.shuffle(all_clients)
        return all_clients[:self.config['federated']['clients_per_round']]
    
    def _setup_logging(self):
        """配置日志记录"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
            handlers=[
                logging.FileHandler('federated.log'),
                logging.StreamHandler()
            ]
        )
        self.logger = logging.getLogger('FederatedConfig')
        
    def get_model_config(self) -> Dict[str, Any]:
        """获取模型配置"""
        return self.config['federated']['model']
    
    def get_optimizer_config(self) -> Dict[str, Any]:
        """获取优化器配置"""
        return self.config['federated']['optimizer']
    
    def get_privacy_config(self) -> Dict[str, Any]:
        """获取隐私保护配置"""
        return self.config['federated']['privacy']
    
    def get_security_config(self) -> Dict[str, Any]:
        """获取安全配置"""
        return self.config['federated']['security']